import time
import pandas as pd
from datetime import datetime

from lightweight_charts import Chart
def rsi(df, period: int = 14):
    return pd.DataFrame(
        {"time": df["date"], "RSI": talib.RSI(df["close"], timeperiod=14)}
    )


def sma(df, period: int = 50):
    return pd.DataFrame(
        {"time": df["time"], f"SMA {period}": df["close"].rolling(window=period).mean()}
    ).dropna()

def main():
    chart = Chart()

    df = pd.read_csv("513500.SH.csv")
    df = df.rename(columns={
        'date': 'time',
        'open': 'open',
        'high': 'high',
        'low': 'low',
        'close': 'close'
        # 如果有成交量列需添加 'volume'
    })
    df['time'] = df['time'].apply(lambda x:datetime.strptime(str(x),'%Y%m%d'))
    df = df.sort_values('time')  # 确保按时间升序排列
    #df.set_index('date', inplace=True)

    print(df[['time', 'open', 'high', 'low', 'close']].head())

    # 保持与 sma 返回的 SMA {period} 相同
    line = chart.create_line("SMA 30")

    rsi_chart = chart.create_subchart(height=0.3, width=1, sync=True)
    rsi_line = rsi_chart.create_line("RSI")
    rsi_line.set(rsi(df))

    line.set(sma(df, period=30))

    chart.set(df)

    chart.show(block=True)


if __name__ == "__main__":
    main()
